#pragma once

#include <queue>
#include <mutex>
#include <condition_variable>
#include <utility>
#include <atomic>

template<class T>
class ThreadSafeQueue
{
public:
	explicit ThreadSafeQueue(size_t capacity = 1000) : mCapacity{ capacity } {}
	// Push an element. Block if queue is full.
	void push(T&& elem);
	// Pop an element and return a copy. Block if queue is empty.
	T pop();
	void SetPushEnable(bool status)
	{
		pushEnable.store(status);
	}
	bool GetPushEnable()
	{
		return pushEnable.load();
	}
	void StopThreadWait()
	{
		stopWait.store(true);
		mNotEmpty.notify_all();
	}
	long GetSize()
	{
		return mQueue.size();
	}
private:
	std::mutex mMutex;
	std::condition_variable mNotEmpty;
	std::condition_variable mNotFull;
	std::queue<T> mQueue;
	const size_t mCapacity;
	std::atomic<bool> pushEnable = false;  //Push使能
	std::atomic<bool> stopWait=false;
};

template<class T>
void ThreadSafeQueue<T>::push(T&& elem)
{
	if (pushEnable.load())
	{
		std::unique_lock<std::mutex> lck(mMutex);
		if (mQueue.size() >= mCapacity) return;
		mNotFull.wait(lck, [this] { return mQueue.size() < mCapacity; });

		mQueue.push(std::forward<T>(elem));
		lck.unlock();

		mNotEmpty.notify_one();
	}
}

template<class T>
T ThreadSafeQueue<T>::pop()
{
	std::unique_lock<std::mutex> lck(mMutex);
	mNotEmpty.wait(lck, [this] { return stopWait.load() || (!mQueue.empty()); });
	if (stopWait.load()) return T();
	T elem = mQueue.front();
	mQueue.pop();
	lck.unlock();

	mNotFull.notify_one();
	return elem;
}